from kubernetes.stream import stream
from langchain.agents import AgentExecutor, create_react_agent
# from langchain.llms import Ollama
from langchain.prompts import PromptTemplate
from langchain_core.callbacks import StreamingStdOutCallbackHandler
from langchain_ollama import OllamaLLM
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
import os


class DocumentStorage:
    def __init__(self, directory_path, index_path="faiss_index"):
        self.directory_path = directory_path
        self.index_path = index_path
        self.embeddings = HuggingFaceEmbeddings(
            model_name=r"model/bge-large-zh",
            model_kwargs={'device': 'cuda'}
        )
        # self.vector_store = self.create_vector_store()
        if os.path.exists(self.index_path):
            self.vector_store = FAISS.load_local(self.index_path, self.embeddings)
        else:
            self.vector_store = self.create_vector_store()

    def create_vector_store(self):
        loader = TextLoader(self.directory_path, encoding="utf-8")
        documents = loader.load()
        vector_store = FAISS.from_documents(documents, self.embeddings)
        vector_store.save_local(self.index_path)
        return vector_store

    def search_document(self, query, topk=5):
        similar_docs = self.vector_store.similarity_search(query, k=topk)
        print(f"相似的文档内容\n")
        for i, doc in enumerate(similar_docs):
            print(f"文档{i + 1}:\n{doc.page_content[:0]}...\n")
        return [doc.page_content for doc in similar_docs]


class ReactAgentSystem:
    def __init__(self, document_storage):
        self.document_storage = document_storage
        self.llm = OllamaLLM(
            model="qwen3:4b",
            temperature=0.3,
            num_ctx=4096,
            stream = True,
            callbacks=[StreamingStdOutCallbackHandler()]
        )

    def create_react_agent(self, query, related_docs):
        react_prompt = PromptTemplate(
            template=(
                "以下是与问题相关的企业文档内容：\n{docs}\n\n"
                "请基于文档信息回答问题：\n问题：{query}\n"
                "如果在文档中没有找到相关内容就直接告诉我没有"
            ),
            input_variables=["query", "docs"]
        )
        return react_prompt.format(query=query, docs="\n".join(related_docs))

    def generate_response(self, query):
        related_docs = self.document_storage.search_document(query)
        prompt = self.create_react_agent(query, related_docs)
        response = self.llm.invoke(prompt)
        print("\nReact Agent回答：\n", response)
        return response


class ZeroShotAgentSystem:
    def __init__(self):
        self.llm = OllamaLLM(
            model="qwen3:4b",
            temperature=0.3,
            num_ctx=4096,
            stream = True,
            callbacks=[StreamingStdOutCallbackHandler()]
        )

    def create_zero_shot_prompt(self, query):
        prompt = PromptTemplate(
            template="请基于常识回答下述问题：\n问题：{query}\n",
            input_variables=["query"]
        )
        return prompt.format(query=query)

    def generate_response(self, query):
        prompt = self.create_zero_shot_prompt(query)
        response = self.llm.invoke(prompt)   # ✅ 修正
        print("\nZero-shot Agent回答：\n", response)
        return response


class EnterpriseAgentSystem:
    def __init__(self, document_directory):
        self.document_storage = DocumentStorage(document_directory)
        self.react_agent = ReactAgentSystem(self.document_storage)
        self.zero_shot_agent = ZeroShotAgentSystem()

    def execute_query(self, query, user_react=True):
        related_docs = self.document_storage.search_document(query)
        if user_react:
            print("使用ReAct Agent进行回答：")
            # 直接调用 react_agent 来生成回答
            response = self.react_agent.generate_response(query)
        else:
            print("使用Zero-shot Agent进行回答：")
            response = self.zero_shot_agent.generate_response(query)

        print(f"最终回答：{response}")
        return response



if __name__ == "__main__":
    # documents_directory = "data/1.txt"
    # enterprise_agent = EnterpriseAgentSystem(documents_directory)
    # query1 = "RK3588淘宝链接"
    # print("查询1结果")
    # enterprise_agent.execute_query(query1, user_react=True)
    #
    # query2 = "申请部门申请日期"
    # print("查询2结果")
    # enterprise_agent.execute_query(query2, user_react=True)
    if __name__ == "__main__":
        # 初始化企业代理系统
        documents_directory = "data/1.txt"
        enterprise_agent = EnterpriseAgentSystem(documents_directory)

        print("欢迎使用企业问答系统！请输入您的问题（输入'退出'结束对话）")

        # 进入问答循环
        while True:
            # 获取用户输入
            user_query = input("\n请输入您的问题: ")

            # 检查是否退出
            if user_query.strip().lower() == "退出":
                print("感谢使用，再见！")
                break

            # 执行查询并获取结果
            print("查询结果:")
            enterprise_agent.execute_query(user_query, user_react=True)